from zhipuai import ZhipuAI
import time
import os
from datetime import datetime
import json
from functools import lru_cache
import requests

# 定义一个名为 SmartSearchAssistant 的类，该类用于实现一个智能搜索助手，
# 能够根据用户的查询判断是否需要联网搜索，生成直接回答或进行搜索并分析结果，
# 同时支持调用工具（如获取天气信息）来辅助回答。
class SmartSearchAssistant:
    def __init__(self):
        self.client = ZhipuAI(api_key=os.getenv("GLM_API_KEY"))
        self.model = "GLM-4-FlashX-250414"
        self.searchModel = "Search-Std"  
        self.conversation_history = []
        # 定义可用工具
        self.tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_current_weather",
                    "description": "获取指定城市的当前天气信息",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "location": {
                                "type": "string",
                                "description": "城市名称(中英文皆可)"
                            }
                        },
                        "required": ["location"]
                    }
                }
            }
        ]
    
    @lru_cache(maxsize=100)
    def determine_search_need(self, query):
        """智能判断是否需要联网搜索"""
        start_time = time.time()
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": """请根据规则判断问题是否需要联网搜索：
    1. 需要搜索的情况：
    - 询问实时信息（股价、天气、赛事）
    - 询问2025年及以后的事件/数据
    - 明确要求"最新"信息

    2. 不需要搜索的情况：
    - 常识性问题
    - 理论/学术问题
    - 2024年及以前的历史事实
    - 操作指南类问题

    请用JSON格式回答：{"need_search": boolean, "reason": string}"""},
                    {"role": "user", "content": query}
                ],
                response_format={"type": "json_object"},
                max_tokens=100
            )
            
            # Fix: Access the message content directly
            content = response.choices[0].message.content
            result = json.loads(content)
            
            return result.get("need_search", False)
            
        except Exception as e:
            print(f"判断失败，默认不搜索：{str(e)}")
            return False
        finally:
            elapsed_time = time.time() - start_time
            print(f"判断是否需要搜索耗时：{elapsed_time:.2f}秒")
    
    def generate_direct_response(self, query):
        """生成直接回答"""
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M")
        self._stream_output(f"\n【{current_time} 知识库回答】\n")
        
        start_time = time.time()
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "你是一个知识丰富的AI助手，请提供准确、专业的回答。"},
                *self._get_conversation_context(),
                {"role": "user", "content": query}
            ],
            stream=True
        )
        
        elapsed_time = time.time() - start_time
        print(f"生成直接回答耗时：{elapsed_time:.2f}秒")
        return self._stream_response(response)
    
    def search_and_analyze(self, query):
        """执行搜索并生成分析"""
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M")
        self._stream_output(f"\n【{current_time} 联网搜索中...】\n")
        
        start_time = time.time()
        try:
            # 执行搜索 - 移除不支持的max_results参数
            search_response = self.client.web_search.web_search(
                search_engine=self.searchModel,
                search_query=query
                # 移除了 max_results=3
            )
            
            if not search_response.search_result:
                raise ValueError("未找到搜索结果")
            
            elapsed_time = time.time() - start_time
            print(f"搜索耗时：{elapsed_time:.2f}秒")
            
            # 显示原始结果
            # self._show_search_results(search_response.search_result)
            
            # 生成分析
            self._stream_output(f"\n【{current_time} 分析总结】\n")
            analysis_response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "请基于以下搜索结果进行专业总结，注明信息时间并指出数据来源"},
                    *self._prepare_search_messages(search_response.search_result),
                    {"role": "user", "content": f"请分析关于'{query}'的搜索结果"}
                ],
                stream=True
            )
            
            return self._stream_response(analysis_response)
            
        except Exception as e:
            print(f"搜索失败：{str(e)}")
            self._stream_output("\n⚠ 搜索失败，转为知识库回答\n")
            return self.generate_direct_response(query)
        finally:
            elapsed_time = time.time() - start_time
            print(f"搜索和分析耗时：{elapsed_time:.2f}秒")
    
    def process_query(self, query):
        """处理用户查询"""
        self.conversation_history.append(("用户", query))
        print("\n助手: ", end="", flush=True)
        
        try:
            start_time = time.time()
            # 让模型自主决定是否需要调用工具
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是一个智能助手，可以根据需要调用工具获取信息。"},
                    *self._get_conversation_context(),
                    {"role": "user", "content": query}
                ],
                tools=self.tools,
                tool_choice="auto",  # 让模型自主决定是否调用工具
            )
            elapsed_time = time.time() - start_time
            print(f"首轮交换耗时：{elapsed_time:.2f}秒")
            
            response_message = response.choices[0].message
            tool_calls = response_message.tool_calls
            
            # 将模型的初始响应（包括可能的工具调用）添加到历史记录
            if response_message.content:
                self.conversation_history.append(("助手", response_message.content))
            
            # 如果有工具调用
            if tool_calls:
                available_functions = {
                    "get_current_weather": self.get_current_weather,
                }
                
                # 处理每个工具调用
                for tool_call in tool_calls:
                    function_name = tool_call.function.name
                    function_to_call = available_functions[function_name]
                    function_args = json.loads(tool_call.function.arguments)
                    
                    # 将工具调用信息添加到历史记录
                    tool_call_msg = f"调用工具: {function_name}({function_args})"
                    self.conversation_history.append(("系统", tool_call_msg))
                    
                    # 调用函数并获取结果
                    function_response = function_to_call(**function_args)
                    
                    # 将工具原始返回结果添加到历史记录
                    self.conversation_history.append(("工具", str(function_response)))
            
            # 如果没有工具调用，按原流程处理
            need_search = not tool_calls and self.determine_search_need(query)
            if need_search:
                response = self.search_and_analyze(query)
            else:
                response = self.generate_direct_response(query)
            
            self.conversation_history.append(("助手", response))
            
            # 总体耗时
            elapsed_time = time.time() - start_time
            print(f"总体耗时：{elapsed_time:.2f}秒")
            return response
            
        except Exception as e:
            error_msg = f"处理错误：{str(e)}"
            self._stream_output(error_msg)
            self.conversation_history.append(("助手", error_msg))
            return error_msg
        
    # 定义天气API的调用函数
    def get_current_weather(self, location):
        start_time = time.time()
        """
        获取指定城市的当前天气信息
        
        参数:
            location (str): 城市名称(中英文皆可)
        
        返回:
            str: 包含天气信息的中文字符串
        """
        # 判断是否为英文，如果不是则调用大模型翻译
        if not location.isascii():
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是一个翻译助手，将中文地名翻译成英文。"},
                    {"role": "user", "content": f"请将以下中文城市名称翻译成英文，只需返回英文名称，不要解释或其他内容：{location}"}
                ],
                temperature=0,
            )
            location = response.choices[0].message.content.strip()
        
        APIkey = 'd94892357e6b825afa1c15556492f40e'  # Your API key
        language = 'zh_cn'  # OpenWeatherMap uses underscore for language codes
        
        url = f'http://api.openweathermap.org/data/2.5/weather?q={location}&units=metric&appid={APIkey}&lang={language}'
        
        try:
            target = requests.get(url)
            content = target.json()
            
            if target.status_code != 200:
                raise ValueError(f"Weather API error: {content.get('message', 'Unknown error')}")
            
            # 组织天气信息为中文字符串
            weather_desc = content["weather"][0]["description"]
            temp = content["main"]["temp"]
            feels_like = content["main"]["feels_like"]
            humidity = content["main"]["humidity"]
            city_name = content.get("name", location)
            
            weather_info = (
                f"{city_name}当前天气情况：\n"
                f"天气状况：{weather_desc}\n"
                f"温度：{temp}°C\n"
                f"体感温度：{feels_like}°C\n"
                f"湿度：{humidity}%"
            )
            
            return weather_info
            
        except Exception as e:
            error_msg = f"查询{city_name}天气信息失败：{str(e)}"
            print(error_msg)
            return error_msg
        finally:
            elapsed_time = time.time() - start_time
            print(f"查询{city_name}天气信息耗时：{elapsed_time:.2f}秒")
            
    # 定义获取城市英文名称的函数
    def getcitynameinenglish(city):
        """
        从本地JSON文件获取城市的英文名称
        
        参数:
            city (str): 中文城市名称
        
        返回:
            str: 对应的英文城市名称，若未找到则返回None
        
        异常:
            FileNotFoundError: 当chinacities.json文件不存在时
            json.JSONDecodeError: 当JSON文件格式错误时
        """
        try:
            with open('chinacities.json', 'r', encoding='utf-8') as file:
                cities_data = json.load(file)
                for item in cities_data:
                    if item['city'] == city:
                        return item['en_name']
        except FileNotFoundError:
            print("File not found.")
        except json.JSONDecodeError:
            print("Error decoding JSON.")
        return None
    
    # 辅助方法 --------------------------------------------------
    def _prepare_search_messages(self, results):
        """准备搜索结果的对话上下文"""
        messages = []
        for i, result in enumerate(results[:3], 1):
            messages.append({
                "role": "system",
                "content": f"搜索结果{i}: {result.title}\n{result.content[:500]}"
            })
        return messages
    
    def _show_search_results(self, results):
        """展示原始搜索结果"""
        self._stream_output("\n🔍 原始搜索结果：\n")
        for i, result in enumerate(results[:3], 1):
            self._stream_output(f"\n▷ 结果{i} ◁\n")
            self._stream_output(f"标题: {result.title}\n")
            content = "\n".join(line for line in result.content.split("\n") if line.strip())
            self._stream_output(f"内容: {content[:300]}...\n")
    
    def _get_conversation_context(self):
        """获取最近的对话上下文"""
        return [
            {"role": "assistant" if role == "助手" else "user", "content": content}
            for role, content in self.conversation_history[-3:]
        ]
    
    def _stream_response(self, response):
        """处理流式响应"""
        full_response = ""
        for chunk in response:
            if chunk.choices and chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                print(content, end="", flush=True)
                full_response += content
                time.sleep(0.02)
        return full_response
    
    def _stream_output(self, text, speed=0.03):
        """通用流式输出"""
        for char in text:
            print(char, end="", flush=True)
            time.sleep(speed)
    
    def show_history(self):
        """显示对话历史"""
        print("\n=== 对话历史 ===")
        for role, text in self.conversation_history:
            print(f"{role}: {text}\n")

# 主程序 ------------------------------------------------------
if __name__ == "__main__":
    assistant = SmartSearchAssistant()
    print("智能助手已启动（输入'退出'结束）\n")
    
    while True:
        try:
            user_input = input("您的问题：\n> ").strip()
            
            if user_input.lower() in ["退出", "exit", "quit"]:
                print("\n对话结束")
                break
                
            if not user_input:
                continue
                
            # 显示loading提示
            print("\n助手: 正在思考中...", end="", flush=True)
            assistant.process_query(user_input)
            
            print()  # 空行分隔
            
        except KeyboardInterrupt:
            print("\n对话中断")
            break
